import numpy as np
import pandas as pd

from AsmExcel import AsmExcel
from AutogluonModel import AutogluonModel
from Fy3dL1Hdf import Fy3dL1Hdf
from Fy3dL2VsmHdf import Fy3dL2VsmHdf
from GeoTiff import GeoTiff
from Timer import Timer
from utils import list_files, find_files


def l1_contains_hn_data(filepath, date, hour):
    # 风云卫星经过河南的可能时间段：升轨13-15点、降轨01-03点
    # l1文件名时间格式  yyyymmdd_hh
    l1_hour_str = f"{date}_{hour}"
    return (filepath.upper().endswith('.HDF') and
            l1_hour_str in filepath and
            (("ASCEND" in filepath and hour in ["13", "14", "15"]) or
             "DESCEND" in filepath and hour in ["01", "02", "03"]))


def vsm_contains_hn_data(filepath, date, hour):
    # vsm文件名时间格式  yyyymmdd
    return (filepath.upper().endswith('.HDF') and
            date in filepath and
            hour in ["01", "02", "03", "13", "14", "15"])


class Asm10Model(AutogluonModel):

    def __init__(self, path):
        AutogluonModel.__init__(self, path)
        self.target = AsmExcel.VWC

    def prepare_train_data(self, asm_dir, l1_dir=None, vsm_dir=None, tvdi_dir=None, ndvi_dir=None, lst_dir=None, soiltype_file=None, landuse_file=None, save_path=None):
        timer = Timer()
        timer.tick(f"开始准备训练数据")
        if soiltype_file:
            timer.tick(f"加载soiltype文件：{soiltype_file}")
            soiltype_tiff = GeoTiff(soiltype_file)
        if landuse_file:
            timer.tick(f"加载landuse文件：{landuse_file}")
            landuse_tiff = GeoTiff(landuse_file)
        # 找到ASM10米的文件
        asm10_paths = list_files(asm_dir, lambda f: '_10.xls' in f.lower())
        combined_data = []
        for asm10_path in asm10_paths:
            timer.tick(f"开始处理asm文件：{asm10_path}")
            asm10_xls = AsmExcel(asm10_path)
            # 提取入库时间 yyyy_mm_dd_hh
            hour_strs = asm10_xls.append_col_timestr("%Y_%m_%d_%H", "hour_str")
            for hour_str in hour_strs:
                timer.tick(f"开始处理小时数据：{hour_str}")
                hour_str_parts = hour_str.split("_")
                date = f"{hour_str_parts[0]}{hour_str_parts[1]}{hour_str_parts[2]}"
                hour = hour_str_parts[3]
                if l1_dir:
                    l1_paths = find_files(l1_dir, lambda f: l1_contains_hn_data(f, date, hour))
                    if len(l1_paths) == 0:
                        timer.tick(f"未匹配到合理的l1文件，略过处理本小时数据。")
                        continue
                if vsm_dir:
                    vsm_paths = list_files(vsm_dir, lambda f: vsm_contains_hn_data(f, date, hour))
                    if len(vsm_paths) == 0:
                        timer.tick(f"未匹配到合理的vsm文件，略过处理本小时数据。")
                        continue
                if tvdi_dir:
                    tvdi_paths = list_files(tvdi_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
                    if len(tvdi_paths) == 0:
                        timer.tick(f"未匹配到tvdi文件，略过处理本小时数据。")
                        continue
                if ndvi_dir:
                    ndvi_paths = list_files(ndvi_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
                    if len(ndvi_paths) == 0:
                        timer.tick(f"未匹配到ndvi文件，略过处理本小时数据。")
                        continue
                if lst_dir:
                    lst_paths = list_files(lst_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
                    if len(lst_paths) == 0:
                        timer.tick(f"未匹配到lst文件，略过处理本小时数据。")
                        continue
                # 同时找到匹配的l1、ndvi和lst文件，开始构建记录
                if l1_dir:
                    l1_path = l1_paths[0]
                    timer.tick(f"匹配到l1文件：{l1_path}")
                    l1_hdf = Fy3dL1Hdf(l1_path)
                    timer.tick(f"开始l1文件地图映射（经纬度范围：{l1_hdf.latlon_area}）：{l1_path}  ")
                    l1_hdf.map_to_area()
                    timer.tick(f"l1文件地图映射完成：{l1_path}")
                if vsm_dir:
                    vsm_path = vsm_paths[0]
                    timer.tick(f"匹配到vsm文件：{vsm_path}")
                    vsm_hdf = Fy3dL2VsmHdf(vsm_path)
                if tvdi_dir:
                    tvdi_path = tvdi_paths[0]
                    timer.tick(f"匹配到tvdi文件：{tvdi_path}")
                    tvdi_tiff = GeoTiff(tvdi_path)
                if ndvi_dir:
                    ndvi_path = ndvi_paths[0]
                    timer.tick(f"匹配到ndvi文件：{ndvi_path}")
                    ndvi_tiff = GeoTiff(ndvi_path)
                if lst_dir:
                    lst_path = lst_paths[0]
                    timer.tick(f"匹配到lst文件：{lst_path}")
                    lst_tiff = GeoTiff(lst_path)
                rows = asm10_xls.sub_rows(lambda r: r["hour_str"] == hour_str)
                n = 0
                for row in rows.itertuples(index=False, name=None):
                    station_no, lat, lon, height, city, station_name, county, timestamp, vmc, srh, gwc, aws, _ = row
                    combined_row = [station_no, station_name, lon, lat, hour_str, vmc, srh, gwc, aws]
                    if l1_dir and l1_hdf.in_area(lat, lon):
                        earth_ob_data = l1_hdf.earth_ob_at(lat, lon)
                        combined_row.extend(earth_ob_data)
                    if vsm_dir:
                        if hour in ["01", "02", "03"]:
                            combined_row.append(vsm_hdf.vsm_a_at(lat, lon))
                        elif hour in ["13", "14", "15"]:
                            combined_row.append(vsm_hdf.vsm_d_at(lat, lon))
                        else:
                            combined_row.append(float('nan'))
                    if tvdi_dir:
                        combined_row.append(tvdi_tiff.get_pixel_value_interpolated(lat, lon))
                    if ndvi_dir:
                        combined_row.append(ndvi_tiff.get_pixel_value_interpolated(lat, lon))
                    if lst_dir:
                        combined_row.append(lst_tiff.get_pixel_value_interpolated(lat, lon))
                    if soiltype_file:
                        combined_row.append(soiltype_tiff.get_pixel_value(lat, lon))
                    if landuse_file:
                        combined_row.append(landuse_tiff.get_pixel_value(lat, lon))
                    combined_data.append(combined_row)
                    n += 1
                    if timer.elapsed(60):
                        timer.tick(f"正在处理小时数据：{hour_str}，已完成{n}条，共{len(rows)}条")
                timer.tick(f"小时数据处理完成：{hour_str}，共{len(rows)}条")
            timer.tick(f"asm文件处理完成：{asm10_path}")

        timer.tick(f"生成训练数据集")
        columns = [AsmExcel.STATION_NO, AsmExcel.STATION_NAME, AsmExcel.LONGITUDE, AsmExcel.LATITUDE, '观测时间',
                   AsmExcel.VWC, AsmExcel.SRH, AsmExcel.GWC, AsmExcel.AWS]
        features = []
        if l1_dir:
            features.extend(['earth_ob_1', 'earth_ob_2', 'earth_ob_3', 'earth_ob_4', 'earth_ob_5',
                            'earth_ob_6', 'earth_ob_7', 'earth_ob_8', 'earth_ob_9', 'earth_ob_10'])
        if vsm_dir:
            features.append("vsm")
        if tvdi_dir:
            features.append("tvdi")
        if ndvi_dir:
            features.append("ndvi")
        if lst_dir:
            features.append("lst")
        if soiltype_file:
            features.append("soiltype")
        if landuse_file:
            features.append("landuse")
        columns.extend(features)
        combined_df = pd.DataFrame(combined_data, columns=columns)
        # 立即转换分类列
        if soiltype_file:
            combined_df["soiltype"] = combined_df["soiltype"].astype('category')
        if landuse_file:
            combined_df["landuse"] = combined_df["landuse"].astype('category')
        timer.tick(f"准备训练数据完成")
        self.df = combined_df
        self.features = features
        if save_path:
            combined_df.to_excel(save_path, index=False)
        return combined_df

    def prepare_data(self, hour_str, l1_dir=None, vsm_dir=None, tvdi_dir=None, ndvi_dir=None, lst_dir=None, soiltype_file=None, landuse_file=None, save_path=None):
        if not tvdi_dir and not ndvi_dir and not lst_dir:
            raise ValueError("需要 tvdi_dir/ndvi_dir/lst_dir 其中的至少一个。")
        timer = Timer()
        timer.tick(f"开始准备数据")
        if soiltype_file:
            timer.tick(f"加载soiltype文件：{soiltype_file}")
            soiltype_tiff = GeoTiff(soiltype_file)
        if landuse_file:
            timer.tick(f"加载landuse文件：{landuse_file}")
            landuse_tiff = GeoTiff(landuse_file)
        combined_data = []
        timer.tick(f"开始处理小时数据：{hour_str}")
        hour_str_parts = hour_str.split("_")
        date = f"{hour_str_parts[0]}{hour_str_parts[1]}{hour_str_parts[2]}"
        hour = hour_str_parts[3]
        if l1_dir:
            l1_paths = find_files(l1_dir, lambda f: l1_contains_hn_data(f, date, hour))
            if len(l1_paths) == 0:
                timer.tick(f"未匹配到合理的l1文件。")
                return
        if vsm_dir:
            vsm_paths = list_files(vsm_dir, lambda f: vsm_contains_hn_data(f, date, hour))
            if len(vsm_paths) == 0:
                timer.tick(f"未匹配到合理的vsm文件。")
                return
        if tvdi_dir:
            tvdi_paths = list_files(tvdi_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
            if len(tvdi_paths) == 0:
                timer.tick(f"未匹配到tvdi文件。")
                return
        if ndvi_dir:
            ndvi_paths = list_files(ndvi_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
            if len(ndvi_paths) == 0:
                timer.tick(f"未匹配到ndvi文件。")
                return
        if lst_dir:
            lst_paths = list_files(lst_dir, lambda f: f.upper().endswith('.TIF') and hour_str in f)
            if len(lst_paths) == 0:
                timer.tick(f"未匹配到lst文件。")
                return
        # 同时找到匹配的l1、ndvi和lst文件，开始构建记录
        sample_tiff = None
        if l1_dir:
            l1_path = l1_paths[0]
            timer.tick(f"匹配到l1文件：{l1_path}")
            l1_hdf = Fy3dL1Hdf(l1_path)
            timer.tick(f"开始l1文件地图映射（经纬度范围：{l1_hdf.latlon_area}）：{l1_path}  ")
            l1_hdf.map_to_area()
            timer.tick(f"l1文件地图映射完成：{l1_path}")
        if vsm_dir:
            vsm_path = vsm_paths[0]
            timer.tick(f"匹配到vsm文件：{vsm_path}")
            vsm_hdf = Fy3dL2VsmHdf(vsm_path)
        if tvdi_dir:
            tvdi_path = tvdi_paths[0]
            timer.tick(f"匹配到tvdi文件：{tvdi_path}")
            tvdi_tiff = GeoTiff(tvdi_path)
            sample_tiff = tvdi_tiff
        if ndvi_dir:
            ndvi_path = ndvi_paths[0]
            timer.tick(f"匹配到ndvi文件：{ndvi_path}")
            ndvi_tiff = GeoTiff(ndvi_path)
            sample_tiff = ndvi_tiff
        if lst_dir:
            lst_path = lst_paths[0]
            timer.tick(f"匹配到lst文件：{lst_path}")
            lst_tiff = GeoTiff(lst_path)
            sample_tiff = lst_tiff
        if not sample_tiff:
            raise ValueError("需要 tvdi_dir/ndvi_dir/lst_dir 其中的至少一个。")
        n = 0
        # 按照sample_tiff的每一个像素进行循环
        for col in range(0, sample_tiff.cols):
            for row in range(0, sample_tiff.rows):
                lat, lon = sample_tiff.get_latlon(col, row)
                combined_row = [col, row, lon, lat]
                if l1_dir and l1_hdf.in_area(lat, lon):
                    earth_ob_data = l1_hdf.earth_ob_at(lat, lon)
                    combined_row.extend(earth_ob_data)
                if vsm_dir:
                    if hour in ["01", "02", "03"]:
                        combined_row.append(vsm_hdf.vsm_a_at(lat, lon))
                    elif hour in ["13", "14", "15"]:
                        combined_row.append(vsm_hdf.vsm_d_at(lat, lon))
                    else:
                        combined_row.append(float('nan'))
                if tvdi_dir:
                    combined_row.append(tvdi_tiff.get_value(col, row))
                if ndvi_dir:
                    combined_row.append(ndvi_tiff.get_value(col, row))
                if lst_dir:
                    combined_row.append(lst_tiff.get_value(col, row))
                if soiltype_file:
                    combined_row.append(soiltype_tiff.get_pixel_value(lat, lon))
                if landuse_file:
                    combined_row.append(landuse_tiff.get_pixel_value(lat, lon))
                if all([~np.isnan(v) for v in combined_row]):
                    combined_data.append(combined_row)
                n += 1
                if timer.elapsed(60):
                    timer.tick(f"正在处理小时数据：{hour_str}，已完成{n}条，共{sample_tiff.cols * sample_tiff.rows}条")
        timer.tick(f"小时数据处理完成：{hour_str}，共{sample_tiff.cols * sample_tiff.rows}条")

        timer.tick(f"生成数据集")
        columns = ["col", "row", "lon", "lat"]
        features = []
        if l1_dir:
            features.extend(['earth_ob_1', 'earth_ob_2', 'earth_ob_3', 'earth_ob_4', 'earth_ob_5',
                            'earth_ob_6', 'earth_ob_7', 'earth_ob_8', 'earth_ob_9', 'earth_ob_10'])
        if vsm_dir:
            features.append("vsm")
        if tvdi_dir:
            features.append("tvdi")
        if ndvi_dir:
            features.append("ndvi")
        if lst_dir:
            features.append("lst")
        if soiltype_file:
            features.append("soiltype")
        if landuse_file:
            features.append("landuse")
        columns.extend(features)
        combined_df = pd.DataFrame(combined_data, columns=columns)
        # 立即转换分类列
        if soiltype_file:
            combined_df["soiltype"] = combined_df["soiltype"].astype('category')
        if landuse_file:
            combined_df["landuse"] = combined_df["landuse"].astype('category')
        timer.tick(f"准备数据完成")
        self.df = combined_df
        self.features = features
        self.sample_tiff = sample_tiff
        if save_path:
            combined_df.to_excel(save_path, index=False)
        return combined_df

    def export_tiffs(self, result_df, sample_tiff=None, save_path_func=None):
        if sample_tiff is None:
            sample_tiff = self.sample_tiff
        if sample_tiff is None:
            raise ValueError("未指定样例tiff，需要参与本次运算的 tvdi/ndvi/lst 其中的至少一个tiff文件。")
        timer = Timer()
        # 计算偏移量索引
        cols = result_df["col"].astype(int)
        rows = result_df["row"].astype(int)
        for label in self.predictor.labels:
            save_path = save_path_func(label)
            timer.tick(f"开始导出{label}结果tiff文件：{save_path}")
            array = np.full((sample_tiff.rows, sample_tiff.cols), sample_tiff.nodata_value, dtype=np.float32)
            array[rows, cols] = result_df[label].values
            sample_tiff.save_as(save_path, array)
            timer.tick(f"{save_path} 导出完成。")
